"""
风险管理模块
提供全面的风险控制和评估功能
"""

import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional, Union


class RiskManager:
    """风险管理器"""
    
    def __init__(self, initial_capital: float = 100000.0):
        """
        初始化风险管理器
        
        Args:
            initial_capital: 初始资金
        """ 
        self.initial_capital = initial_capital
        self.current_capital = initial_capital
        self.positions = {}  # 持仓记录
        self.trade_history = []  # 交易历史
        self.max_drawdown = 0.0  # 最大回撤
        self.daily_returns = []  # 日收益率
        
    def calculate_position_size(self, entry_price: float, stop_loss: float, 
                              risk_percent: float = 0.02) -> float:
        """
        计算头寸规模
        
        Args:
            entry_price: 入场价格
            stop_loss: 止损价格
            risk_percent: 每笔交易风险百分比
            
        Returns:
            float: 头寸规模（股数）
        """
        if entry_price <= 0 or stop_loss <= 0:
            raise ValueError("价格必须大于0")
            
        # 计算每股风险
        risk_per_share = abs(entry_price - stop_loss)
        
        # 计算总风险金额
        total_risk_amount = self.current_capital * risk_percent
        
        # 计算头寸规模
        position_size = total_risk_amount / risk_per_share
        
        return max(0, position_size)
    
    def calculate_var(self, returns: List[float], confidence_level: float = 0.95) -> float:
        """
        计算风险价值(Value at Risk)
        
        Args:
            returns: 收益率序列
            confidence_level: 置信水平
            
        Returns:
            float: VaR值
        """
        if len(returns) < 2:
            return 0.0
            
        # 计算VaR
        var = np.percentile(returns, (1 - confidence_level) * 100)
        return abs(var)
    
    def calculate_sharpe_ratio(self, returns: List[float], risk_free_rate: float = 0.02) -> float:
        """
        计算夏普比率
        
        Args:
            returns: 收益率序列
            risk_free_rate: 无风险利率
            
        Returns:
            float: 夏普比率
        """
        if len(returns) < 2:
            return 0.0
            
        excess_returns = np.array(returns) - risk_free_rate/252  # 日化无风险利率
        if np.std(excess_returns) == 0:
            return 0.0
            
        sharpe = np.mean(excess_returns) / np.std(excess_returns) * np.sqrt(252)  # 年化
        return sharpe
    
    def update_position(self, symbol: str, quantity: float, price: float, 
                       trade_type: str, timestamp: datetime = None):
        """
        更新持仓
        
        Args:
            symbol: 股票代码
            quantity: 数量
            price: 价格
            trade_type: 交易类型 ('buy' 或 'sell')
            timestamp: 时间戳
        """
        if timestamp is None:
            timestamp = datetime.now()
            
        # 记录交易
        trade_record = {
            'timestamp': timestamp,
            'symbol': symbol,
            'quantity': quantity,
            'price': price,
            'trade_type': trade_type,
            'value': quantity * price
        }
        self.trade_history.append(trade_record)
        
        # 更新持仓
        if trade_type == 'buy':
            if symbol in self.positions:
                # 加仓
                current_qty = self.positions[symbol]['quantity']
                current_avg_price = self.positions[symbol]['avg_price']
                new_qty = current_qty + quantity
                new_avg_price = (current_qty * current_avg_price + quantity * price) / new_qty
                
                self.positions[symbol] = {
                    'quantity': new_qty,
                    'avg_price': new_avg_price,
                    'last_update': timestamp
                }
            else:
                # 新建仓位
                self.positions[symbol] = {
                    'quantity': quantity,
                    'avg_price': price,
                    'last_update': timestamp
                }
        elif trade_type == 'sell':
            if symbol in self.positions:
                current_qty = self.positions[symbol]['quantity']
                if quantity >= current_qty:
                    # 清仓
                    del self.positions[symbol]
                else:
                    # 减仓
                    self.positions[symbol]['quantity'] -= quantity
                    self.positions[symbol]['last_update'] = timestamp
    
    def calculate_portfolio_risk(self) -> Dict[str, float]:
        """
        计算投资组合风险指标
        
        Returns:
            Dict: 风险指标字典
        """
        # 计算最大回撤
        if len(self.daily_returns) > 1:
            cumulative_returns = np.cumsum(self.daily_returns)
            running_max = np.maximum.accumulate(cumulative_returns)
            drawdown = cumulative_returns - running_max
            max_drawdown = abs(np.min(drawdown))
            self.max_drawdown = max(self.max_drawdown, max_drawdown)
        
        return {
            'current_capital': self.current_capital,
            'max_drawdown': self.max_drawdown,
            'total_trades': len(self.trade_history),
            'positions_count': len(self.positions)
        }
    
    def get_risk_report(self) -> str:
        """
        生成风险报告
        
        Returns:
            str: 风险报告文本
        """
        risk_metrics = self.calculate_portfolio_risk()
        
        report = f"""
风险控制报告
{'='*30}
当前资金: ¥{risk_metrics['current_capital']:,.2f}
最大回撤: {risk_metrics['max_drawdown']:.2%}
交易次数: {risk_metrics['total_trades']}
持仓数量: {risk_metrics['positions_count']}
"""
        
        return report


def main():
    """主函数 - 风险管理演示"""
    print("风险管理模块演示")
    print("="*30)
    
    # 创建风险管理器
    rm = RiskManager(initial_capital=100000.0)
    
    # 示例：计算头寸规模
    entry_price = 100.0
    stop_loss = 95.0
    position_size = rm.calculate_position_size(entry_price, stop_loss, risk_percent=0.02)
    print(f"\n头寸规模计算示例:")
    print(f"入场价: ¥{entry_price}")
    print(f"止损价: ¥{stop_loss}")
    print(f"头寸规模: {position_size:.0f} 股")
    
    # 示例：更新持仓
    rm.update_position("AAPL", 100, 150.0, "buy")
    rm.update_position("TSLA", 50, 200.0, "buy")
    
    # 生成风险报告
    print(rm.get_risk_report())

if __name__ == "__main__":
    main()